{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "ba98beac-d461-4d7d-878a-11beca03ea1c",
      "metadata": {
        "id": "ba98beac-d461-4d7d-878a-11beca03ea1c"
      },
      "source": [
        "# Time travel\n",
        "\n",
        "## Review\n",
        "\n",
        "We discussed motivations for human-in-the-loop:\n",
        "\n",
        "(1) `Approval` - We can interrupt our agent, surface state to a user, and allow the user to accept an action\n",
        "\n",
        "(2) `Debugging` - We can rewind the graph to reproduce or avoid issues\n",
        "\n",
        "(3) `Editing` - You can modify the state\n",
        "\n",
        "We showed how breakpoints can stop the graph at specific nodes or allow the graph to dynamically interrupt itself.\n",
        "\n",
        "Then we showed how to proceed with human approval or directly edit the graph state with human feedback.\n",
        "\n",
        "## Goals\n",
        "\n",
        "Now, let's show how LangGraph [supports debugging](https://langchain-ai.github.io/langgraph/how-tos/human_in_the_loop/time-travel/) by viewing, re-playing, and even forking from past states.\n",
        "\n",
        "We call this `time travel`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "bd48aeb6-8478-4cb4-aef1-d524b80824d3",
      "metadata": {
        "id": "bd48aeb6-8478-4cb4-aef1-d524b80824d3"
      },
      "outputs": [],
      "source": [
        "%%capture --no-stderr\n",
        "%pip install --quiet -U langgraph langchain_google_genai langgraph_sdk"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "7d32093f",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7d32093f",
        "outputId": "ff2899ad-66a8-41cd-8099-ead5db6837fb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "env: GOOGLE_API_KEY=AIzaSyDojJx56O8nzFauBKc_rj_0o7wvAnQgHAI\n",
            "env: LANGCHAIN_API_KEY=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\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "from google.colab import userdata\n",
        "\n",
        "%env GOOGLE_API_KEY = {userdata.get('GEMINI_API_KEY')}\n",
        "\n",
        "%env LANGCHAIN_API_KEY = {userdata.get('LANGCHAIN_API_KEY')}\n",
        "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
        "os.environ[\"LANGCHAIN_PROJECT\"] = \"langchain-academy\""
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0497d316-832a-4668-b133-fd317ee81220",
      "metadata": {
        "id": "0497d316-832a-4668-b133-fd317ee81220"
      },
      "source": [
        "Let's build our agent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "d64ab3a1-b39c-4176-88c7-791a0b80c725",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d64ab3a1-b39c-4176-88c7-791a0b80c725",
        "outputId": "194c8255-cbfc-4da6-cdcf-956b3223e0bc"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:langchain_google_genai._function_utils:Key 'title' is not supported in schema, ignoring\n",
            "WARNING:langchain_google_genai._function_utils:Key 'title' is not supported in schema, ignoring\n",
            "WARNING:langchain_google_genai._function_utils:Key 'title' is not supported in schema, ignoring\n",
            "WARNING:langchain_google_genai._function_utils:Key 'title' is not supported in schema, ignoring\n",
            "WARNING:langchain_google_genai._function_utils:Key 'title' is not supported in schema, ignoring\n",
            "WARNING:langchain_google_genai._function_utils:Key 'title' is not supported in schema, ignoring\n",
            "WARNING:langchain_google_genai._function_utils:Key 'title' is not supported in schema, ignoring\n",
            "WARNING:langchain_google_genai._function_utils:Key 'title' is not supported in schema, ignoring\n",
            "WARNING:langchain_google_genai._function_utils:Key 'title' is not supported in schema, ignoring\n"
          ]
        }
      ],
      "source": [
        "from langchain_google_genai import ChatGoogleGenerativeAI\n",
        "from langchain_core.tools import tool\n",
        "\n",
        "def multiply(a: int, b: int) -> int:\n",
        "    \"\"\"Multiply a and b.\n",
        "\n",
        "    Args:\n",
        "        a: first int\n",
        "        b: second int\n",
        "    \"\"\"\n",
        "    return a * b\n",
        "\n",
        "# This will be a tool\n",
        "def add(a: int, b: int) -> int:\n",
        "    \"\"\"Adds a and b.\n",
        "\n",
        "    Args:\n",
        "        a: first int\n",
        "        b: second int\n",
        "    \"\"\"\n",
        "    return a + b\n",
        "\n",
        "def divide(a: int, b: int) -> float:\n",
        "    \"\"\"Adds a and b.\n",
        "\n",
        "    Args:\n",
        "        a: first int\n",
        "        b: second int\n",
        "    \"\"\"\n",
        "    return a / b\n",
        "\n",
        "tools: list[tool] = [add, multiply, divide]\n",
        "llm: ChatGoogleGenerativeAI = ChatGoogleGenerativeAI(model = \"gemini-1.5-flash\")\n",
        "llm_with_tools = llm.bind_tools(tools)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from langgraph.prebuilt import tools_condition, ToolNode\n",
        "\n",
        "ToolNode"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 204
        },
        "id": "8VpqJbqJBsEy",
        "outputId": "54f9f0a2-e3c0-4fc0-951e-76fbdde819c2"
      },
      "id": "8VpqJbqJBsEy",
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "langgraph.prebuilt.tool_node.ToolNode"
            ],
            "text/html": [
              "<div style=\"max-width:800px; border: 1px solid var(--colab-border-color);\"><style>\n",
              "      pre.function-repr-contents {\n",
              "        overflow-x: auto;\n",
              "        padding: 8px 12px;\n",
              "        max-height: 500px;\n",
              "      }\n",
              "\n",
              "      pre.function-repr-contents.function-repr-contents-collapsed {\n",
              "        cursor: pointer;\n",
              "        max-height: 100px;\n",
              "      }\n",
              "    </style>\n",
              "    <pre style=\"white-space: initial; background:\n",
              "         var(--colab-secondary-surface-color); padding: 8px 12px;\n",
              "         border-bottom: 1px solid var(--colab-border-color);\"><b>langgraph.prebuilt.tool_node.ToolNode</b><br/>def __init__(tools: Sequence[Union[BaseTool, Callable]], *, name: str=&#x27;tools&#x27;, tags: Optional[list[str]]=None, handle_tool_errors: Optional[bool]=True) -&gt; None</pre><pre class=\"function-repr-contents function-repr-contents-collapsed\" style=\"\"><a class=\"filepath\" style=\"display:none\" href=\"#\">/usr/local/lib/python3.10/dist-packages/langgraph/prebuilt/tool_node.py</a>A node that runs the tools called in the last AIMessage.\n",
              "\n",
              "It can be used either in StateGraph with a &quot;messages&quot; key or in MessageGraph. If\n",
              "multiple tool calls are requested, they will be run in parallel. The output will be\n",
              "a list of ToolMessages, one for each tool call.\n",
              "\n",
              "The `ToolNode` is roughly analogous to:\n",
              "\n",
              "```python\n",
              "tools_by_name = {tool.name: tool for tool in tools}\n",
              "def tool_node(state: dict):\n",
              "    result = []\n",
              "    for tool_call in state[&quot;messages&quot;][-1].tool_calls:\n",
              "        tool = tools_by_name[tool_call[&quot;name&quot;]]\n",
              "        observation = tool.invoke(tool_call[&quot;args&quot;])\n",
              "        result.append(ToolMessage(content=observation, tool_call_id=tool_call[&quot;id&quot;]))\n",
              "    return {&quot;messages&quot;: result}\n",
              "```\n",
              "\n",
              "Important:\n",
              "    - The state MUST contain a list of messages.\n",
              "    - The last message MUST be an `AIMessage`.\n",
              "    - The `AIMessage` MUST have `tool_calls` populated.</pre>\n",
              "      <script>\n",
              "      if (google.colab.kernel.accessAllowed && google.colab.files && google.colab.files.view) {\n",
              "        for (const element of document.querySelectorAll('.filepath')) {\n",
              "          element.style.display = 'block'\n",
              "          element.onclick = (event) => {\n",
              "            event.preventDefault();\n",
              "            event.stopPropagation();\n",
              "            google.colab.files.view(element.textContent, 67);\n",
              "          };\n",
              "        }\n",
              "      }\n",
              "      for (const element of document.querySelectorAll('.function-repr-contents')) {\n",
              "        element.onclick = (event) => {\n",
              "          event.preventDefault();\n",
              "          event.stopPropagation();\n",
              "          element.classList.toggle('function-repr-contents-collapsed');\n",
              "        };\n",
              "      }\n",
              "      </script>\n",
              "      </div>"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "1d8622a9-57cd-44dc-8696-46c5ab32d0b9",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 266
        },
        "id": "1d8622a9-57cd-44dc-8696-46c5ab32d0b9",
        "outputId": "00249a9b-8b61-43d8-db0b-642ef0e514de"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/jpeg": "/9j/4AAQSkZJRgABAQAAAQABAAD/4gHYSUNDX1BST0ZJTEUAAQEAAAHIAAAAAAQwAABtbnRyUkdCIFhZWiAH4AABAAEAAAAAAABhY3NwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQAA9tYAAQAAAADTLQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAlkZXNjAAAA8AAAACRyWFlaAAABFAAAABRnWFlaAAABKAAAABRiWFlaAAABPAAAABR3dHB0AAABUAAAABRyVFJDAAABZAAAAChnVFJDAAABZAAAAChiVFJDAAABZAAAAChjcHJ0AAABjAAAADxtbHVjAAAAAAAAAAEAAAAMZW5VUwAAAAgAAAAcAHMAUgBHAEJYWVogAAAAAAAAb6IAADj1AAADkFhZWiAAAAAAAABimQAAt4UAABjaWFlaIAAAAAAAACSgAAAPhAAAts9YWVogAAAAAAAA9tYAAQAAAADTLXBhcmEAAAAAAAQAAAACZmYAAPKnAAANWQAAE9AAAApbAAAAAAAAAABtbHVjAAAAAAAAAAEAAAAMZW5VUwAAACAAAAAcAEcAbwBvAGcAbABlACAASQBuAGMALgAgADIAMAAxADb/2wBDAAMCAgMCAgMDAwMEAwMEBQgFBQQEBQoHBwYIDAoMDAsKCwsNDhIQDQ4RDgsLEBYQERMUFRUVDA8XGBYUGBIUFRT/2wBDAQMEBAUEBQkFBQkUDQsNFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBT/wAARCAD5ANYDASIAAhEBAxEB/8QAHQABAAIDAQEBAQAAAAAAAAAAAAUGAwQHCAECCf/EAFEQAAEEAQIDAgYLDAcGBwAAAAEAAgMEBQYRBxIhEzEVFiJBUZQIFBcyVVZhdNHS0yM1NlRxdYGRk5WytCU3QkNSgpIYJGRylqEzNFNiscHw/8QAGwEBAQADAQEBAAAAAAAAAAAAAAECAwUEBgf/xAAzEQEAAQIBCQUJAQADAAAAAAAAAQIRAwQSITFBUVKR0RQzYXGhBRMVI2KSscHhgSLw8f/aAAwDAQACEQMRAD8A/qmiIgIiICIiAsNq5XpR89ieOuz/ABSvDR+sqDu37uevz47FTGlVrnkt5NrQ5zX/APpQhwLS4d7nuBa3cNAc4u5Ptbh/p+F5llxcF+ydua1fb7ZmcR5y9+5/V0W+KKae8n/IW29u+NWF+F6HrLPpTxqwvwxQ9ZZ9KeKuF+B6HqzPoTxVwvwPQ9WZ9CvyfH0XQeNWF+GKHrLPpTxqwvwxQ9ZZ9KeKuF+B6HqzPoTxVwvwPQ9WZ9CfJ8fQ0HjVhfhih6yz6U8asL8MUPWWfSnirhfgeh6sz6E8VcL8D0PVmfQnyfH0NB41YX4Yoess+lblTIVb7S6rZhstHeYZA4D9S0/FXC/A9D1Zn0LUtaB05bkErsNThnad22K0QhmafkkZs4foKfJnbPp/E0J9FWI7NzSM8MN+1NksPK4RsvT8va1XE7NbKQAHMPQB+24O3NvuXCzrXXRm+MEwIiLWgiIgIiICIiAiIgIiICIiAojV2Yfp/S+VyMQDpq1Z8kTXdxft5IP6dlLqvcQqct7ROZjhaZJm13SsY0blzmeWAB6SW7LbgxE4lMVarwsa0hp/Dx4DDVKEZ5uxZ5cnnkkJ3e8/K5xc4n0kqRWGnaivVILMDueGZjZGO9LSNwf1FZlhVMzVM1a0FUuIHFbS3C6LHv1JkzSfkJHRVIIa01madzW8z+SKFj3kNHUnbYbjchW1cU9krQqPg07k48frBupMc+zJiM5o7HG7NQldG0OZNEA4Ojl6Atc0tPL1LehWI2cp7JjT+N4q6b0m2tetUc3hfC8OTq463ODzyQthaGxwu8lzZHOdISAzZodylwVgtcftBUdct0hZz3tfOvtNotilpzthNhw3bCJzH2XaHcbN59zuBsuUx5fWendd8Ltfax0nlrtuxpGzicxDp6g+4+neklrTDnij3LWu7J43G4aehPnVA4t4/Wep5tTDMYbX+W1Bj9VwW8fUxsEwwsOJguRSRyRtjIjsSGJpJGz5ec9GgDoHpi3x20TT1je0ocpYsahozR17VCnjbVh8DpI2yMLzHE4NYWvb5ZPLuSN9wQIvgLx7xvHPBWblWjdx1yvYsxyV56VlkYjZYkijc2aSJjHuc1gc5jSSwktcAQtbhLp+7jOMXGnJWsbYqQZLLY91W3NA5jbUbMdA0ljiNnta/nb03APMO/dRfsY7GQ0vh8poTMaezWNyWLymUte3rFF7aFmGW9JLG6GxtyPLmzNPKDuOV24GyDuCIiDXyFCvlaFmlbibPVsxuhlif3PY4bOB/KCVEaGvz39Nwi1L29upLNRmlO+8j4ZXRF53/wAXJzfpU+qzw8b2mn5Lg35L921cj5htvHJO90Z2+VnKf0r0U9zVffH7XYsyIi86CIiAiIgIiICIiAiIgIiICIiCqU52aDeaNvaLAOeXU7fXkqbncwynuY3cnkf0btsw7EN7THqvhFobX+RjyWo9JYTP3mxCFlrIUYp5BGCSGhzgTy7ucdvlKtr2NkY5j2h7HDYtcNwR6Cq0/h9joSTjbOQwoP8AdY62+OIejaI7xt/Q0f8AYL0TVRiaa5tPO/8A3/WWiVePsbeFBaG+5vpblBJA8EwbA+f+z8gVm0fw70tw9hsxaY09jNPxWXNdOzG1GQCUjcAuDQN9tz3+lYfEmx8as9+2h+yTxJsfGrPftofsk93h8fpKWjetCKr+JNj41Z79tD9kqnex2Wr8VcHp5mqcx4OuYW/flJlh7TtYZ6bGbfc/e8tiTfp38vUed7vD4/SS0b3VFC6s0XgNd4xuO1HhaGdx7ZBM2rka7Z4w8AgO5XAjcBxG/wApWj4k2PjVnv20P2SeJNj41Z79tD9knu8Pj9JLRvQDfY3cKWBwbw40u0PGzgMTB1G4Ox8n0gfqUnpngroDRmXiyuA0XgcNk4g5sdyjj4oZWhw2cA5rQRuCQVueJNj41Z79tD9kvviBTsO/pDIZXKs337G1deIj+VjOVrh8jgQmZhxrr5R/4Wh+crkPG7t8Nipeeo/mhyGRhd5ELOodFG4d8p7unvBu4kHla6ywQR1oI4YWNiijaGMYwbBrQNgAPMF8q1YaVeOvXhjrwRtDWRRNDWtA7gAOgCyrCuuJjNp1QSIiLUgiIgIiICIiAiIgIiICIiAiIgIiICIiAufZYt937SwJPN4sZfYebb21jd/P+TzfpHn6Cuf5Xf3ftLdW7eLGX6EDf/zWN7vPt+Tp3b+ZB0BERAREQEREBERAREQEREBERAREQEREBERAREQEREBERAXPcsB/tA6VPM0HxXzHk7dT/veM677d36fOP0dCXPctt/tBaV6nm8V8xsOX/i8Z5/8A9/2QdCREQEREBERAREQEREBERAREQEREBERAREQERaeXy1fB46a7aLhDEBuGNLnOJIDWtA7ySQAPOSFYiaptGsbiKlP1Dquby4cVia7HdRHYuyOkaP8A3cse2/pAJHylfnw7rD8Qwfrc32a9fZa98c4Wy7oqR4d1h+IYP1ub7NPDusPxDB+tzfZp2WvfHOCy7rwHrH2e2V097IivibXCud2ocTHc06MfFmA7t5Z7FZzXsd7X35T7XG2w8oPB8wXsXw7rD8Qwfrc32a5BnvY/zah9kHh+LVjH4YZnHVexNQWJDFPM0csU7j2e/Oxp2H/Kz/D1dlr3xzgs9LIqR4d1h+IYP1ub7NPDusPxDB+tzfZp2WvfHOCy7oqR4d1h+IYP1ub7NPDusPxDB+tzfZp2WvfHOCy7oqUzPaua7d+NwsjR3tbdmaT+nsjt+pWPAZyHP0PbEbHwSMeYpq8u3PDI33zHbdOnpG4IIIJBBWqvArw4zp1eE3LJJERaEEREBERAREQEREBERAREQFUuJh2wVEeY5ahuD85jVtVR4m/eKh+dqH8zGvTk3f0ecMqdcNtERepiIiICKJy2qsXgsthsbesmG7mJn16MXZvd2r2RukcNwCG7Ma47uIHTbv6KRt24KFWazZmjr1oWOklmlcGsY0DcucT0AAG5JUGVFr43I1cxjqt+lPHapWomTwTxO5mSRuAc1zT5wQQR+VbCoItXKZWng8bayORtQ0aFWJ009mw8MjijaN3Oc49AAASSVmrzx2oI5oXiSKRoex7e5zSNwQgyLR0Af6V1kPMMszYAf8DVK3lo6A++2s/zvH/I1VZ7uvy/cMo1SuKIi5bEREQEREBERAREQEREBERAVR4m/eKh+dqH8zGrcqjxN+8VD87UP5mNenJu/o84ZU64bapHGvU1PSPDDOZG7NlIIuSOux2EkbHddLLI2KJsTndGuc97W8x6DffzK7qK1TpbFa10/dwecpR5HFXGdnPWl32eNwR1BBBBAIIIIIBBBC9M6mLzLpStxPZkOJvD+nl7uIzEunamSw5y2ddl5aU0kk0bh7adG1zecRjps4MPVpO6zxRal1Nw/vYHS1rWdfUWAz1eTUun8rqD+k3V3QbmCpf3I5H7tla7mbzbOG7AQF2Cn7HXh9RZkBFgXOfkaRx92aW/ZkltQl7X8skjpC55BY3lc4lzQNmkAkL432OfD5mAfhm4KVtR91uQfK3I2hadYawxtkNjte1JDCWjd/QEha82RzHEanhzWq+BGS03qLU8mOyFrK421WzN6UvkMNS04stRc3LJJHKzbmIJ8huzj0Kr+GqZbFaV17ozXuX1W/XE+mb1508uZfNjclCwnexU5SDAQSxrotmbNdts4EleiMZwk0jhYdMQ0MNHUi00+aTFMhlkaK75Y3xyu995Zc2R+5fzHdxPf1WnovgZofh9dtW8HgmVrFisab3z2JrPLXJ5jCwSvcGRk7Esbs07Dp0VzZHG8fVq6V9jvw0wuOvatv5fVUdD2lXx+oJYZnymkJHsFmQuNes1jHOLY9tthyjqVWqWqtaxcOsjp/IagymPyWN4lY7AMuw5Q27UVSZ9ZzojZdG0zbdu8cz2dRsCDsu8wexv4eVdPHBw4KWPGCyy5FE3JWg6tKwODHQP7Xmg2D3DaMtGziNtlt47gHoLEV5IKWAbWhkyFTKvjjtThr7dYh0M5HP1eCAXE+/I8vmUzZHCeKWPt4vTXsgdFSZ3OZLCUtKVszT9v5OaeeCR7LPaR9s5xe6JxgYSxxLdi4bbOIXonhXputpfQmIq1bmQvRSV45+1yV+W5Ju5jTsHyucQ30NB2HmC3Z9AaftZjN5SfGxz3M1RjxuQdK5z2WKzO05Y3MJ5dvusm+wBPN136L8aD4eYHhnhXYnTtSaljzJ2vYy25rHKeVrdmmV7i1oa1oDQQBt0CyiLSLGtHQH321n+d4/5Gqt5aOgPvtrP87x/yNVbJ7uvy/cMo1SuKIi5bEREQEREBERAREQEREBERAVR4m/eKh+dqH8zGrcorU2D8YcPLTbN7WmD45oZuXm7OWN4ewkbjcczRuNxuNxuN1vwKooxaaqtUTCxoloooZ9/UVfyJdJ2rEg6OfSuVnRH5WmSRjtvytB+RanjPmDfbTbo3LvmLXOcWTVHMZy8m4e8TcrXESNIaSCRuQCGkjoZn1R90dSyyIoTwtnviZlfWqX26eFs98TMr61S+3TM+qPujqtk2ihPC2e+JmV9apfbqr3eMdbH8Qsfoexg78WqshUfdrY4z1eaSFm/M7m7blHc47E7kNJA2BTM+qPujqWdDRQnhbPfEzK+tUvt08LZ74mZX1ql9umZ9UfdHUsm0UJ4Wz3xMyvrVL7dPC2e+JmV9apfbpmfVH3R1LJtaOgPvtrP87x/yNVRGP1RlcpI+GHSmRgsNBJiuWK0TmgPczmLe1Lw0ljtncpDgNwSCFbdKYObC0rDrcrJb92c2rJi37Nry1rQ1m/Xla1jW7nbfbfYb7DXiTFGHVEzGnRomJ2xOzyNUJtERcxiIiICIiAiIgIiICIiAiIgIvjnBjS5xDWgbknuCgY32NT2GyRyTUsRBOfeiNzcpGYuhDtyWxczz3crnOiBB7M/dA/M+Qs6lE1bEyy06ZjhlZnIuykilBk8uOEbkl3I07vLeUdowt5yHBstjcVTw8MkNGrFUikmksPbEwNDpJHl8jzt3uc5xJPnJKzVq0NKtFXrxMggiYI44omhrWNA2DQB0AA6bLKgIiIC/njxB9jLxuz3suqmsq2otK1c/OZszi43XbRigqVJYIhA8iv5xYjBABB3fufT/Q5c/wAhyzcfMByhpdX0zkec7nmaJLVHl6d2x7J3+n8qDoCIiAiIgis3p2vmWPla99DJivJWr5WqyP21Va8tLuzc9rhtzMjcWuBa4sbzNcBstV+opcRekhzcUNKpLahq0L0cjntsukb0bIOUdi/nBYASWu5o9ncz+Rs+iAirIqy6Jqh1NktrT9WCxNNWHbWrjHc3aNEI3c57QC9oiAJADGsGwDVYoJ47MLJoniSJ7Q5rm9xB7igyIiICIiAiIgIiICIiAiLFan9q1ppuR8vZsL+SMbudsN9gPOUEBZEOsr1zHu5J8JUdJTyVK5j+eO690bHBjXv8l0bQ883K1wL9m8wMcjDZFA6Dj5NF4R3a5SYyVI5i/Nn/AH3d7Q4iYDoHjm2LR0BGw6AKeQEREBERAXPuHBOq9Q6g1xvzUciIsdiHb7h9GAvInHXbaWWWZwI99G2E+jb96ltS8QsrY0pjJnR4iu8Mz+Qhc5ruXYO9pROHdI8Edo4Hdkbths+RrmXqvXiqQRwQRshhiaGMjjaGtY0DYAAdwA8yDIiIgIiICIiAoG7RfgbdrK0Ws7CeT2xkoXNlke8Nj5eeJrOby+VrByhp5+UDoepnkQa2OyNXMY+rfo2I7dK1E2eCxC4OZLG4BzXNI6EEEEH5Vsqv4WWSjqTMYuR+UtMcGZGGzbiBrxtlLmmvFKO8sdEXlrurRMzYkbBtgQEREBERAREQERQuY1tp7T9oVsnnMdj7JHN2Nm0xj9vTyk77LOmiqubUxeVtdNIqt7qWjvjTiPXY/pVZ4l3+G3FfQmZ0ln9R4qbFZSDsZQy/G17SCHMe07++a9rXDfpu0bgjotvZ8bgnlK5s7kjoXiBpeGWpow6k31NSdLSGKzuQidmJxCXDtnx83O8PjYJWv28qNzXnvKvy/nF7CngvR4K+yJ1ff1Hm8XJj8PTNbE5T2ywRXDM4fdIzvtuI2uDh3tL9j8vvT3UtHfGnEeux/SnZ8bgnlJmzuWlFVvdS0d8acR67H9Ke6lo7404j12P6U7PjcE8pM2dy0qm57O5DUGXk05puXsJIi0ZXM8vM3HsI37KLccr7Lm9zTuImuEjwd445ojJcRqus86zS+ls5UgfLHz28vFPG50LCPeVmu3Esx9OxZGOrtzysdesHg6Gm8XDjsbWbVpw8xbG0kkuc4ue9zjuXOc5znOc4lznOJJJJK1VUVUTauLJaz5gcDQ0xiK2MxlcVqVcEMZzFxJJLnOc5xLnvc4lznuJc5ziSSSSpBEWCCIiAiIgIiICIiCu2yG8Q8UN8yS/F3OkX3tHLNW/8b0Tnm+5+lgn9CsS45k/ZFcKq/EbFQy8T8LE9mNvtfEzO1Bjw4TVBtP8AdOk469mP8Ptj0LsaAiIgIiICIiDSzVx2Pw960wAvggklaD6WtJH/AMKo6SqR1sBSkA5p7MTJ55ndXzSOaC57iepJJ/R3dwVn1V+DGY+ZzfwFV7TX4OYr5pF/AF0MDRhT5rsSSIizQREQEREGrksbWy1OStajEkT/AJdi0jqHNI6tcDsQ4dQQCOq39B5SfNaLwd60/tbM9OJ8sm23O7lG7tvNueu3yrEsPCz+rnTnzGL+FY4unBnwmPxPRdi0oiLnIIiICIq3rrWcGisQLDoxZuTv7KrV5uXtX95JPma0bkn0DYbkgHZh4dWLXFFEXmRM5PLUcJUdbyNyvQqt99PalbGwflc4gKsS8YdHQvLTnIXEdN445Hj9YaQuH5O1azuR8IZWw6/e68skg8mIb+9jb3Mb0HQdTsCST1WNfW4XsPDin5tc38P7cvDuPuzaN+Gm+ry/UT3ZtG/DTfV5fqLhyLd8Dybiq5x0Lw4FxI9jppPVPsxsdqSvcjPD3JSeGMq4RSBsdhh3fBy7c33V/Keg2Ae70L3d7s2jfhpvq8v1Fw5E+B5NxVc46F4dx92bRvw031eX6i+s4yaNe7bw3G35XwyNH6y1cNRPgeTcVXOOheHpbD6gxmoa7p8XkKuQiaeVzq0rZA0+g7HofkKkF5YgMlK9HepTyUb8fvLVchr2/IehDh0HkuBB26gruvDfXw1jSmr22sgy9MNE8bPeytPdKweZpIII72kEdRsTxcu9l1ZLT7yib0+sLr1LkiIuEiL1V+DGY+ZzfwFV7TX4OYr5pF/AFYdVfgxmPmc38BVe01+DmK+aRfwBdHB7mfP9Lsb1h0jIJHQsbLMGksY53KHO26AnY7dfPsV524W8etUYzgrmNZ68xUVivUvW4Ks2Puiazdn8ISV46wh7GNrNnckbXcx5gOYhvVejV57h4Baul0DqXQU+RwsWAdfmy+By0Jldchsm8LkTZ4i0M5WvLmkteSRt0Ck32IsDfZCT6WtZmpxD0wdIWqGFlz8XtXINyEdmtE4Nla14YzaVrnMHJtsecbOIWCvxvzs9iriNT6Om0dNqDF27WEsx5Ntpz3xQ9q6KUNY0wyhh5wAXDyXeVuFG5ngRqji5kM3e4i3MNRdPp2xp+hU086WaOHt3NdJZe+VrCXbxx7MA2AB3J71u47hRrrV+qtNZHX9/BMqaap2oajMCZnvuWJ4DXdPL2jWiMCMv2Y3m6vPldAp/yEHpLjjmNNcMOC2MixbtV6o1XhGTNnyuWFRkj4oInSc072vL5XmQbN2Jds4kjZehMfNPZoVprNY07MkTXy1y8P7J5AJZzDodjuNx0Oy8/WOC2vncEMDw9sUdC6ir4+pJjpJMr7ZaOzY1rKtiPlY4smaA4uA8+3K8Ltmg9P29KaJwGFv5KTMXsdQgqT5CbfnsvZGGukO5J3cQT1JPXqSrTfaJ1YeFn9XOnPmMX8KzLDws/q5058xi/hVxe5nzj8SuxaURFzkEREBcC4s5J2S4iWIHOJixtWOCNp7muk+6PI/KOyB/5Au+rgXFnGuxnEOedzSIsnVjnjee5z4/ubwPyDsj/nC73sXN7Vp12m3p+rrslVkWvkb8WLoz25xKYYWF7xDC+V+w9DGAucfkAJVVHFvT5/us5/07kPsF9vViUUaKpiGtcnODWkkgAdST5lxOl7KDD3chUeyDHnCW7bKkU7M1A695T+RsjqY8sMLiD74uDTuWhXtnFHT997avY5o9uez2fp++xp36dXGAADr3k7KvcPtCau0HFj9Ptfp+9pmhI5sV6Zsovur7ktYWAcnMNwOfm7h73deTErrrqp9zVo22tO637Vin43X68OUyUmli3T2LzMmHuX/CDe0aW2BCJWRcnlN3c0kFzSNyBzAbnX4mcUMxNh9c0dL4Sa5BhaM8V3NNvisas5gL9oRsS98bXNcdi3Y9Ad1nyPCbL2+HWsMAyzSFzMZ2bJ13ue/s2xPtsmAeeTcO5WkbAEb+fzrBqHhprCv484/TlnCyYTVQmmkGTdMyarYlgEUhbyNIe13K09dtj6fPoqnKM2030x4X2/wdH0XPLa0dgpppHzTSUIHvkkcXOc4xtJJJ7yT51MKi4/W+K0bjKGDvtykl3H1oa0zqeFvTxFzY2glsjIS1w+UFZ/dd08f7rO/9O5D7Be2nFw4iImqL+aLmpbRWSdh9e4CyxxaJpzSlA/tslaQB/rEbv8qreFzVbP46O7UFhsDyQBarS15Oh2O7JGtcO7zjqrJonGuzOvcBWY3mbBObspH9hkbSQf8AWYx/mUyiaJwK5q1Wn8Mqdb0giIvzBUXqr8GMx8zm/gKr2mvwcxXzSL+AK05mm7I4i9UYQHzwSRAnzFzSP/tVDSVyOxgacIPJZrQsgsQO6Phka0BzHA9QQf1jYjoQuhgacKY8V2JhERZoIiICIiAsPCz+rnTnzGL+FY8nlK2IqPs2pRHG3oB3ue49A1rR1c4kgBo3JJAHUqQ0Ji58JozCUbTOzswU4mSx778j+Ubt38+x6b/IscXRgz4zH4nquxOoiLnIIiICrmudGQa1w4rPkFa3C/tatrl5jE/u6jpu0jcEb9x6EEAixotmHiVYVcV0TaYHl3K1LWn8h7Qy1c4+515WvO7JR/ijf3PHd3dRuNw09FjXpzJYulmaj6t+pBerP99DZibIw/laQQqxLwg0dK4uOBrtJ67RuewfqBAX1uF7cw5p+bRN/D+locKRdy9xvRvwHF+1k+snuN6N+A4v2sn1lu+OZNw1co6locNRdy9xvRvwHF+1k+snuN6N+A4v2sn1k+OZNw1co6locNRdy9xvRvwHF+1k+svrODujWO38BQO+R73uH6i7ZPjmTcNXKOpaN7hdYS5C8yjRgkv33+9q1wHPPynrs0dR5TiAN+pXduHGgho2jNPaeyfL2+UzyM95G0e9iYe8tBJO56uJJ2A2a2xYjBY3AVzBjKFbHwk7llaJsYcfSdh1Pylb64mXe1Ksrp93RFqfWV1ahERcNBQuY0Vp/UNgWMpg8bkZwOUS2qkcjwPRu4E7KaRZU11UTembSalW9yvRnxTwn7vi+qnuV6M+KeE/d8X1VaUW7tGNxzzlbzvVb3K9GfFPCfu+L6qe5Xoz4p4T93xfVVpRO0Y3HPOS871W9yvRnxTwn7vi+qnuV6M+KeE/d8X1VaUTtGNxzzkvO9B4rQ2nMFZbZx2AxlCw3flmrVI43t379iBuN1OIi1VV1VzeqbprERFgCIiAiIgIiICIiAiIgIiICIiAiIg//9k=\n",
            "text/plain": [
              "<IPython.core.display.Image object>"
            ]
          },
          "metadata": {}
        }
      ],
      "source": [
        "from IPython.display import Image, display\n",
        "\n",
        "from langgraph.checkpoint.memory import MemorySaver\n",
        "from langgraph.graph import MessagesState\n",
        "from langgraph.graph import START, END, StateGraph\n",
        "from langgraph.prebuilt import tools_condition, ToolNode\n",
        "from langgraph.graph.state import CompiledStateGraph\n",
        "\n",
        "from langchain_core.messages import AIMessage, HumanMessage, SystemMessage\n",
        "\n",
        "# System message\n",
        "sys_msg = SystemMessage(content=\"You are a helpful assistant tasked with performing arithmetic on a set of inputs.\")\n",
        "\n",
        "# Node\n",
        "def assistant(state: MessagesState):\n",
        "   return {\"messages\": [llm_with_tools.invoke([sys_msg] + state[\"messages\"])]}\n",
        "\n",
        "# Graph\n",
        "builder: StateGraph = StateGraph(MessagesState)\n",
        "\n",
        "# Define nodes: these do the work\n",
        "builder.add_node(\"assistant\", assistant)\n",
        "builder.add_node(\"tools\", ToolNode(tools))\n",
        "\n",
        "# Define edges: these determine the control flow\n",
        "builder.add_edge(START, \"assistant\")\n",
        "builder.add_conditional_edges(\n",
        "    \"assistant\",\n",
        "    # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools\n",
        "    # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END\n",
        "    tools_condition,\n",
        ")\n",
        "builder.add_edge(\"tools\", \"assistant\")\n",
        "\n",
        "memory: MemorySaver = MemorySaver()\n",
        "graph: CompiledStateGraph = builder.compile(checkpointer=MemorySaver())\n",
        "\n",
        "# Show\n",
        "display(Image(graph.get_graph(xray=True).draw_mermaid_png()))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fab18a04-1329-47ac-a25b-4e01bf756e2a",
      "metadata": {
        "id": "fab18a04-1329-47ac-a25b-4e01bf756e2a"
      },
      "source": [
        "Let's run it, as before."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "id": "05b2ab62-82bc-4356-8d5b-2d4f49069fdd",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "05b2ab62-82bc-4356-8d5b-2d4f49069fdd",
        "outputId": "91c60284-081c-4c97-9d6a-93e65b1d942a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "================================\u001b[1m Human Message \u001b[0m=================================\n",
            "\n",
            "Multiply 2 and 3\n",
            "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
            "Tool Calls:\n",
            "  multiply (b4ef056c-4332-43da-9a1d-997e5b0ac764)\n",
            " Call ID: b4ef056c-4332-43da-9a1d-997e5b0ac764\n",
            "  Args:\n",
            "    a: 2.0\n",
            "    b: 3.0\n",
            "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
            "Name: multiply\n",
            "\n",
            "6\n",
            "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
            "\n",
            "The result is 6.\n"
          ]
        }
      ],
      "source": [
        "# Input\n",
        "initial_input = {\"messages\": HumanMessage(content=\"Multiply 2 and 3\")}\n",
        "\n",
        "# Thread\n",
        "thread = {\"configurable\": {\"thread_id\": \"1\"}}\n",
        "\n",
        "# Run the graph until the first interruption\n",
        "for event in graph.stream(initial_input, thread, stream_mode=\"values\"):\n",
        "    event['messages'][-1].pretty_print()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "268cfa43-22d1-4d63-8d81-a3ce00f1f2c8",
      "metadata": {
        "id": "268cfa43-22d1-4d63-8d81-a3ce00f1f2c8"
      },
      "source": [
        "## Browsing History\n",
        "\n",
        "We can use `get_state` to look at the **current** state of our graph, given the `thread_id`!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "id": "161eb053-18f6-4c99-8674-8cbd11cae57e",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "161eb053-18f6-4c99-8674-8cbd11cae57e",
        "outputId": "1c23085f-0f34-48a7-f7a2-fe275940f8c6"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "StateSnapshot(values={'messages': [HumanMessage(content='Multiply 2 and 3', additional_kwargs={}, response_metadata={}, id='0e3a5b6c-c2b0-4f39-9350-c09843c7a92e'), AIMessage(content='', additional_kwargs={'function_call': {'name': 'multiply', 'arguments': '{\"b\": 3.0, \"a\": 2.0}'}}, response_metadata={'prompt_feedback': {'block_reason': 0, 'safety_ratings': []}, 'finish_reason': 'STOP', 'safety_ratings': [{'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HARASSMENT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HATE_SPEECH', 'probability': 'NEGLIGIBLE', 'blocked': False}]}, id='run-1d18029a-b2f1-4a64-9ec6-87b61ff7ae04-0', tool_calls=[{'name': 'multiply', 'args': {'b': 3.0, 'a': 2.0}, 'id': '2a3ecfee-d1f2-41db-b3d7-966da5a6ad4c', 'type': 'tool_call'}], usage_metadata={'input_tokens': 208, 'output_tokens': 18, 'total_tokens': 226}), ToolMessage(content='6', name='multiply', id='6c61d4b6-9fa5-43dd-b186-3f02b0f026c7', tool_call_id='2a3ecfee-d1f2-41db-b3d7-966da5a6ad4c'), AIMessage(content='The result is 6. \\n', additional_kwargs={}, response_metadata={'prompt_feedback': {'block_reason': 0, 'safety_ratings': []}, 'finish_reason': 'STOP', 'safety_ratings': [{'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HATE_SPEECH', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HARASSMENT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', 'probability': 'NEGLIGIBLE', 'blocked': False}]}, id='run-b92b3af5-b031-486e-975d-4d3cbb6f6f7a-0', usage_metadata={'input_tokens': 241, 'output_tokens': 6, 'total_tokens': 247}), HumanMessage(content='Multiply 2 and 3', additional_kwargs={}, response_metadata={}, id='9f0994ba-5011-4367-8749-bdcb097349d0'), AIMessage(content='', additional_kwargs={'function_call': {'name': 'multiply', 'arguments': '{\"a\": 2.0, \"b\": 3.0}'}}, response_metadata={'prompt_feedback': {'block_reason': 0, 'safety_ratings': []}, 'finish_reason': 'STOP', 'safety_ratings': [{'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HARASSMENT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HATE_SPEECH', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT', 'probability': 'NEGLIGIBLE', 'blocked': False}]}, id='run-cad7a153-8f25-4f31-a618-701acaae04b5-0', tool_calls=[{'name': 'multiply', 'args': {'a': 2.0, 'b': 3.0}, 'id': 'b4ef056c-4332-43da-9a1d-997e5b0ac764', 'type': 'tool_call'}], usage_metadata={'input_tokens': 257, 'output_tokens': 18, 'total_tokens': 275}), ToolMessage(content='6', name='multiply', id='f094d8c4-b1a1-4e39-b73d-8373a4ba7de3', tool_call_id='b4ef056c-4332-43da-9a1d-997e5b0ac764'), AIMessage(content='The result is 6. \\n', additional_kwargs={}, response_metadata={'prompt_feedback': {'block_reason': 0, 'safety_ratings': []}, 'finish_reason': 'STOP', 'safety_ratings': [{'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HATE_SPEECH', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HARASSMENT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', 'probability': 'NEGLIGIBLE', 'blocked': False}]}, id='run-c5226e1c-c9f2-436b-8d22-718ef6c563d0-0', usage_metadata={'input_tokens': 290, 'output_tokens': 6, 'total_tokens': 296})]}, next=(), config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef7bfb9-5fda-612d-8008-e0df3c5dd3bd'}}, metadata={'source': 'loop', 'writes': {'assistant': {'messages': [AIMessage(content='The result is 6. \\n', additional_kwargs={}, response_metadata={'prompt_feedback': {'block_reason': 0, 'safety_ratings': []}, 'finish_reason': 'STOP', 'safety_ratings': [{'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HATE_SPEECH', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HARASSMENT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', 'probability': 'NEGLIGIBLE', 'blocked': False}]}, id='run-c5226e1c-c9f2-436b-8d22-718ef6c563d0-0', usage_metadata={'input_tokens': 290, 'output_tokens': 6, 'total_tokens': 296})]}}, 'step': 8, 'parents': {}}, created_at='2024-09-26T11:36:35.520930+00:00', parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef7bfb9-5b51-6a83-8007-e317553eda42'}}, tasks=())"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ],
      "source": [
        "graph.get_state({'configurable': {'thread_id': '1'}})"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8d00869e-7b41-4d71-ad3c-cacf8f9c029f",
      "metadata": {
        "id": "8d00869e-7b41-4d71-ad3c-cacf8f9c029f"
      },
      "source": [
        "We can also browse the state history of our agent.\n",
        "\n",
        "`get_state_history` lets us get the state at all prior steps.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "id": "3010169c-3bfa-498c-a30c-7ba53744e4d5",
      "metadata": {
        "id": "3010169c-3bfa-498c-a30c-7ba53744e4d5"
      },
      "outputs": [],
      "source": [
        "all_states = [s for s in graph.get_state_history(thread)]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "id": "c4612ccf-59fc-4848-8845-0433fee2ca8e",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c4612ccf-59fc-4848-8845-0433fee2ca8e",
        "outputId": "951f8909-e803-4cdb-e0a8-234342b56a02"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "10"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ],
      "source": [
        "len(all_states)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "af30f269-1152-4fa1-a7c6-2947acad9a27",
      "metadata": {
        "id": "af30f269-1152-4fa1-a7c6-2947acad9a27"
      },
      "source": [
        "The first element is the current state, just as we got from `get_state`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "id": "4e60b292-8efc-4cc3-b836-51f060fa608b",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4e60b292-8efc-4cc3-b836-51f060fa608b",
        "outputId": "14cd38f5-941c-4f1b-ef0a-c4a2c2f3b20d"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "StateSnapshot(values={'messages': [HumanMessage(content='Multiply 2 and 3', additional_kwargs={}, response_metadata={}, id='0e3a5b6c-c2b0-4f39-9350-c09843c7a92e')]}, next=('assistant',), config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef7bf9d-c6c7-6d9c-8000-19c59752e49f'}}, metadata={'source': 'loop', 'writes': None, 'step': 0, 'parents': {}}, created_at='2024-09-26T11:24:14.694502+00:00', parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef7bf9d-c6c0-64b7-bfff-10ee07d43da6'}}, tasks=(PregelTask(id='b46a6f34-f9f0-5de5-6fea-5d3e373ba386', name='assistant', path=('__pregel_pull', 'assistant'), error=None, interrupts=(), state=None),))"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ],
      "source": [
        "all_states[-2]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4148a710-ceed-413b-b93c-070c6c792fa2",
      "metadata": {
        "id": "4148a710-ceed-413b-b93c-070c6c792fa2"
      },
      "source": [
        "Everything above we can visualize here:\n",
        "\n",
        "![fig1.jpg](https://cdn.prod.website-files.com/65b8cd72835ceeacd4449a53/66dbb038211b544898570be3_time-travel1.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a5ad554a-faf3-489f-a9a9-774f4ec2a526",
      "metadata": {
        "id": "a5ad554a-faf3-489f-a9a9-774f4ec2a526"
      },
      "source": [
        "## Replaying\n",
        "\n",
        "We can re-run our agent from any of the prior steps.\n",
        "\n",
        "![fig2.jpg](https://cdn.prod.website-files.com/65b8cd72835ceeacd4449a53/66dbb038a0bd34b541c78fb8_time-travel2.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e135d2db-d613-42da-877e-d429f21aaefd",
      "metadata": {
        "id": "e135d2db-d613-42da-877e-d429f21aaefd"
      },
      "source": [
        "Let's look back at the step that recieved human input!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "id": "3688e511-a440-4330-a450-e5ed889c3b30",
      "metadata": {
        "id": "3688e511-a440-4330-a450-e5ed889c3b30"
      },
      "outputs": [],
      "source": [
        "to_replay = all_states[-2]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "id": "72adf296-d519-4bdc-af03-3b29799e9534",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "72adf296-d519-4bdc-af03-3b29799e9534",
        "outputId": "20c7c513-e29b-46ad-c696-20c61dac5222"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "StateSnapshot(values={'messages': [HumanMessage(content='Multiply 2 and 3', additional_kwargs={}, response_metadata={}, id='0e3a5b6c-c2b0-4f39-9350-c09843c7a92e')]}, next=('assistant',), config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef7bf9d-c6c7-6d9c-8000-19c59752e49f'}}, metadata={'source': 'loop', 'writes': None, 'step': 0, 'parents': {}}, created_at='2024-09-26T11:24:14.694502+00:00', parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef7bf9d-c6c0-64b7-bfff-10ee07d43da6'}}, tasks=(PregelTask(id='b46a6f34-f9f0-5de5-6fea-5d3e373ba386', name='assistant', path=('__pregel_pull', 'assistant'), error=None, interrupts=(), state=None),))"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ],
      "source": [
        "to_replay"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "571e7894-6546-48ff-9c25-fa6d120391b3",
      "metadata": {
        "id": "571e7894-6546-48ff-9c25-fa6d120391b3"
      },
      "source": [
        "Look at the state."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "id": "6fe69428-f364-4330-bf5d-aa966c7f3b07",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6fe69428-f364-4330-bf5d-aa966c7f3b07",
        "outputId": "aa7adc00-9ec9-49dc-aa00-31ccf624da0a"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'messages': [HumanMessage(content='Multiply 2 and 3', additional_kwargs={}, response_metadata={}, id='0e3a5b6c-c2b0-4f39-9350-c09843c7a92e')]}"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ],
      "source": [
        "to_replay.values"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ff2df545-cc80-4962-a34a-faac7af8eb3d",
      "metadata": {
        "id": "ff2df545-cc80-4962-a34a-faac7af8eb3d"
      },
      "source": [
        "We can see the next node to call."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "id": "d2f333f9-9b2b-46f6-ac3a-525f86b20f1b",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d2f333f9-9b2b-46f6-ac3a-525f86b20f1b",
        "outputId": "a288143b-1cd3-4c3c-efca-c8c945ca37ec"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "('assistant',)"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ],
      "source": [
        "to_replay.next"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b8938c18-5c22-47df-b71e-40afa73c87af",
      "metadata": {
        "id": "b8938c18-5c22-47df-b71e-40afa73c87af"
      },
      "source": [
        "We also get the config, which tells us the `checkpoint_id` as well as the `thread_id`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "id": "b1298786-afa5-4277-927e-708a8629231b",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "b1298786-afa5-4277-927e-708a8629231b",
        "outputId": "05399720-b372-40be-a507-9598fbda6d1b"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'configurable': {'thread_id': '1',\n",
              "  'checkpoint_ns': '',\n",
              "  'checkpoint_id': '1ef7bf9d-c6c7-6d9c-8000-19c59752e49f'}}"
            ]
          },
          "metadata": {},
          "execution_count": 25
        }
      ],
      "source": [
        "to_replay.config"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1d93b5eb-f541-4f82-93b1-48f54bf5cf83",
      "metadata": {
        "id": "1d93b5eb-f541-4f82-93b1-48f54bf5cf83"
      },
      "source": [
        "To replay from here, we simply pass the config back to the agent!\n",
        "\n",
        "The graph knows that this checkpoint has aleady been executed.\n",
        "\n",
        "It just re-plays from this checkpoint!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "id": "531b4cd1-54f6-44aa-9ffe-cf5403dad65d",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "531b4cd1-54f6-44aa-9ffe-cf5403dad65d",
        "outputId": "cbd1730f-47b0-44a7-ab86-51b041e3414b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "================================\u001b[1m Human Message \u001b[0m=================================\n",
            "\n",
            "Multiply 2 and 3\n",
            "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
            "Tool Calls:\n",
            "  multiply (88302836-a4fe-4ab0-a3aa-537f442cffaa)\n",
            " Call ID: 88302836-a4fe-4ab0-a3aa-537f442cffaa\n",
            "  Args:\n",
            "    b: 3.0\n",
            "    a: 2.0\n",
            "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
            "Name: multiply\n",
            "\n",
            "6\n",
            "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
            "\n",
            "The result is 6.\n"
          ]
        }
      ],
      "source": [
        "for event in graph.stream(None, to_replay.config, stream_mode=\"values\"):\n",
        "    event['messages'][-1].pretty_print()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7d7a914e-63e6-4424-970f-15059ce9b4c3",
      "metadata": {
        "id": "7d7a914e-63e6-4424-970f-15059ce9b4c3"
      },
      "source": [
        "Now, we can see our current state after the agent re-ran."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5a5a1f03-19f2-4d22-ba54-1c065ff08e85",
      "metadata": {
        "id": "5a5a1f03-19f2-4d22-ba54-1c065ff08e85"
      },
      "source": [
        "## Forking\n",
        "\n",
        "What if we want to run from that same step, but with a different input.\n",
        "\n",
        "This is forking.\n",
        "\n",
        "![fig3.jpg](https://cdn.prod.website-files.com/65b8cd72835ceeacd4449a53/66dbb038f89f2d847ee5c336_time-travel3.png)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "id": "cdeb5bf2-1566-4d8c-8ea5-65894e3a7038",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cdeb5bf2-1566-4d8c-8ea5-65894e3a7038",
        "outputId": "e2f0738c-7dbb-44bf-a77a-9bb17eeb2646"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[HumanMessage(content='Multiply 2 and 3', additional_kwargs={}, response_metadata={}, id='0e3a5b6c-c2b0-4f39-9350-c09843c7a92e')]"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ],
      "source": [
        "to_fork = all_states[-2]\n",
        "to_fork.values[\"messages\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4a15f6a6-6eaa-48d6-92bb-864ea3a31b6a",
      "metadata": {
        "id": "4a15f6a6-6eaa-48d6-92bb-864ea3a31b6a"
      },
      "source": [
        "Again, we have the config."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "id": "d1621b27-ee51-4dc3-81c4-1d05317280db",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d1621b27-ee51-4dc3-81c4-1d05317280db",
        "outputId": "91855403-ea9d-49af-ebd1-64f70eddfd48"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'configurable': {'thread_id': '1',\n",
              "  'checkpoint_ns': '',\n",
              "  'checkpoint_id': '1ef7bf9d-c6c7-6d9c-8000-19c59752e49f'}}"
            ]
          },
          "metadata": {},
          "execution_count": 29
        }
      ],
      "source": [
        "to_fork.config"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c2102195-0583-4dbe-ad2f-02fac7915585",
      "metadata": {
        "id": "c2102195-0583-4dbe-ad2f-02fac7915585"
      },
      "source": [
        "Let's modify the state at this checkpoint.\n",
        "\n",
        "We can just run `update_state` with the `checkpoint_id` supplied.\n",
        "\n",
        "Remember how our reducer on `messages` works:\n",
        "\n",
        "* It will append, unless we supply a message ID.\n",
        "* We supply the message ID to overwrite the message, rather than appending to state!\n",
        "\n",
        "So, to overwrite the the message, we just supply the message ID, which we have `to_fork.values[\"messages\"].id`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "id": "0b4a918d-858a-41ac-a5d4-e99260e2d6ec",
      "metadata": {
        "id": "0b4a918d-858a-41ac-a5d4-e99260e2d6ec"
      },
      "outputs": [],
      "source": [
        "fork_config = graph.update_state(\n",
        "    to_fork.config,\n",
        "    {\"messages\": [HumanMessage(content='Multiply 5 and 3',\n",
        "                               id=to_fork.values[\"messages\"][0].id)]},\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "id": "8ff4e9bb-8221-42d1-b7d0-b0cbd5dc374a",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8ff4e9bb-8221-42d1-b7d0-b0cbd5dc374a",
        "outputId": "43157fbe-d6c8-49ee-97a3-4bb8bcc5384e"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'configurable': {'thread_id': '1',\n",
              "  'checkpoint_ns': '',\n",
              "  'checkpoint_id': '1ef7bffe-3c3f-61e8-8001-1c17c00e49a8'}}"
            ]
          },
          "metadata": {},
          "execution_count": 31
        }
      ],
      "source": [
        "fork_config"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bebfe6fd-c94b-4291-a125-ec6170e35bc5",
      "metadata": {
        "id": "bebfe6fd-c94b-4291-a125-ec6170e35bc5"
      },
      "source": [
        "This creates a new, forked checkpoint.\n",
        "\n",
        "But, the metadata - e.g., where to go next - is perserved!\n",
        "\n",
        "We can see the current state of our agent has been updated with our fork."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "id": "586ce86c-1257-45e9-ba30-6287932b9484",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "586ce86c-1257-45e9-ba30-6287932b9484",
        "outputId": "4ef838f7-acda-463d-a393-3b09185d4425"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[HumanMessage(content='Multiply 5 and 3', additional_kwargs={}, response_metadata={}, id='0e3a5b6c-c2b0-4f39-9350-c09843c7a92e')]"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ],
      "source": [
        "all_states = [state for state in graph.get_state_history(thread) ]\n",
        "all_states[0].values[\"messages\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "id": "12e19798-25d8-49e8-8542-13d2b3bdf58e",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "12e19798-25d8-49e8-8542-13d2b3bdf58e",
        "outputId": "4515e98a-de4c-422e-b1f6-a2f851aa2b21"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "StateSnapshot(values={'messages': [HumanMessage(content='Multiply 5 and 3', additional_kwargs={}, response_metadata={}, id='0e3a5b6c-c2b0-4f39-9350-c09843c7a92e')]}, next=('assistant',), config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef7bffe-3c3f-61e8-8001-1c17c00e49a8'}}, metadata={'source': 'update', 'step': 1, 'writes': {'__start__': {'messages': [HumanMessage(content='Multiply 5 and 3', additional_kwargs={}, response_metadata={}, id='0e3a5b6c-c2b0-4f39-9350-c09843c7a92e')]}}, 'parents': {}}, created_at='2024-09-26T12:07:23.992076+00:00', parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef7bf9d-c6c7-6d9c-8000-19c59752e49f'}}, tasks=(PregelTask(id='089c07b2-9ec9-a3f4-7fea-3cce02a3ef0a', name='assistant', path=('__pregel_pull', 'assistant'), error=None, interrupts=(), state=None),))"
            ]
          },
          "metadata": {},
          "execution_count": 33
        }
      ],
      "source": [
        "graph.get_state({'configurable': {'thread_id': '1'}})"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "78c641e2-b8e9-4461-b854-8725006a5eb6",
      "metadata": {
        "id": "78c641e2-b8e9-4461-b854-8725006a5eb6"
      },
      "source": [
        "Now, when we stream, the graph knows this checkpoint has never been executed.\n",
        "\n",
        "So, the graph runs, rather than simply re-playing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "id": "1c49f2a8-b325-45e4-b36c-17fab1b37cc0",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1c49f2a8-b325-45e4-b36c-17fab1b37cc0",
        "outputId": "8c31125c-7729-4c94-b9ab-fedd589a8b6c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "================================\u001b[1m Human Message \u001b[0m=================================\n",
            "\n",
            "Multiply 5 and 3\n",
            "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
            "Tool Calls:\n",
            "  multiply (870d99da-b2b8-4766-a062-09036629e9d5)\n",
            " Call ID: 870d99da-b2b8-4766-a062-09036629e9d5\n",
            "  Args:\n",
            "    a: 5.0\n",
            "    b: 3.0\n",
            "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
            "Name: multiply\n",
            "\n",
            "15\n",
            "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
            "\n",
            "The result is 15.\n"
          ]
        }
      ],
      "source": [
        "for event in graph.stream(None, fork_config, stream_mode=\"values\"):\n",
        "    event['messages'][-1].pretty_print()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "428d7f80-ee60-4147-b51f-ee3b0cf5cbba",
      "metadata": {
        "id": "428d7f80-ee60-4147-b51f-ee3b0cf5cbba"
      },
      "source": [
        "Now, we can see the current state is the end of our agent run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "id": "132ef840-64c7-479c-ad34-3f177f4b2524",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "132ef840-64c7-479c-ad34-3f177f4b2524",
        "outputId": "9b6d5702-0552-4f5b-f2fc-b33a363a53e5"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "StateSnapshot(values={'messages': [HumanMessage(content='Multiply 5 and 3', additional_kwargs={}, response_metadata={}, id='0e3a5b6c-c2b0-4f39-9350-c09843c7a92e'), AIMessage(content='', additional_kwargs={'function_call': {'name': 'multiply', 'arguments': '{\"a\": 5.0, \"b\": 3.0}'}}, response_metadata={'prompt_feedback': {'block_reason': 0, 'safety_ratings': []}, 'finish_reason': 'STOP', 'safety_ratings': [{'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HATE_SPEECH', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HARASSMENT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', 'probability': 'NEGLIGIBLE', 'blocked': False}]}, id='run-dc9b607b-81d6-4d54-9fdf-3ceba228ba8a-0', tool_calls=[{'name': 'multiply', 'args': {'a': 5.0, 'b': 3.0}, 'id': '870d99da-b2b8-4766-a062-09036629e9d5', 'type': 'tool_call'}], usage_metadata={'input_tokens': 208, 'output_tokens': 18, 'total_tokens': 226}), ToolMessage(content='15', name='multiply', id='20b4897e-df39-48b1-89a6-d2b5c70189f5', tool_call_id='870d99da-b2b8-4766-a062-09036629e9d5'), AIMessage(content='The result is 15. \\n', additional_kwargs={}, response_metadata={'prompt_feedback': {'block_reason': 0, 'safety_ratings': []}, 'finish_reason': 'STOP', 'safety_ratings': [{'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HATE_SPEECH', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HARASSMENT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', 'probability': 'NEGLIGIBLE', 'blocked': False}]}, id='run-177b1a9f-ee1b-4f51-af26-5430b28e47c8-0', usage_metadata={'input_tokens': 242, 'output_tokens': 7, 'total_tokens': 249})]}, next=(), config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef7c003-e1a4-65d3-8004-fcef10ae8c4b'}}, metadata={'source': 'loop', 'writes': {'assistant': {'messages': [AIMessage(content='The result is 15. \\n', additional_kwargs={}, response_metadata={'prompt_feedback': {'block_reason': 0, 'safety_ratings': []}, 'finish_reason': 'STOP', 'safety_ratings': [{'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HATE_SPEECH', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HARASSMENT', 'probability': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', 'probability': 'NEGLIGIBLE', 'blocked': False}]}, id='run-177b1a9f-ee1b-4f51-af26-5430b28e47c8-0', usage_metadata={'input_tokens': 242, 'output_tokens': 7, 'total_tokens': 249})]}}, 'step': 4, 'parents': {}}, created_at='2024-09-26T12:09:55.552790+00:00', parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef7c003-ddd5-63a8-8003-8b05457c703f'}}, tasks=())"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ],
      "source": [
        "graph.get_state({'configurable': {'thread_id': '1'}})"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2ceb5f31-97b0-466c-9b3b-ae4df7cd462a",
      "metadata": {
        "id": "2ceb5f31-97b0-466c-9b3b-ae4df7cd462a"
      },
      "source": [
        "### Time travel with LangGraph API\n",
        "\n",
        "--\n",
        "\n",
        "**⚠️ DISCLAIMER**\n",
        "\n",
        "*Running Studio currently requires a Mac. If you are not using a Mac, then skip this step.*\n",
        "\n",
        "*Also, if you are running this notebook in CoLab, then skip this step.*\n",
        "\n",
        "--\n",
        "\n",
        "Let's load our `agent` in the Studio UI, which uses `module-3/studio/agent.py` set in `module-3/studio/langgraph.json`.\n",
        "\n",
        "![Screenshot 2024-08-26 at 9.59.19 AM.png](https://cdn.prod.website-files.com/65b8cd72835ceeacd4449a53/66dbb038211b544898570bec_time-travel4.png)\n",
        "\n",
        "We connect to it via the SDK and show how the LangGraph API [supports time travel](https://langchain-ai.github.io/langgraph/cloud/how-tos/human_in_the_loop_time_travel/#initial-invocation)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "891defdb-746c-48e4-8efa-bb5f138dc4bd",
      "metadata": {
        "id": "891defdb-746c-48e4-8efa-bb5f138dc4bd"
      },
      "outputs": [],
      "source": [
        "import platform\n",
        "\n",
        "if 'google.colab' in str(get_ipython()) or platform.system() != 'Darwin':\n",
        "    raise Exception(\"Unfortunately LangGraph Studio is currently not supported on Google Colab or requires a Mac\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a317925d-1788-4cfc-9c12-336b17b4d859",
      "metadata": {
        "id": "a317925d-1788-4cfc-9c12-336b17b4d859"
      },
      "outputs": [],
      "source": [
        "from langgraph_sdk import get_client\n",
        "client = get_client(url=\"http://localhost:62780\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "815d5e03-0ab8-4c7f-a1ee-f410b6aadc03",
      "metadata": {
        "id": "815d5e03-0ab8-4c7f-a1ee-f410b6aadc03"
      },
      "source": [
        "#### Re-playing\n",
        "\n",
        "Let's run our agent streaming `updates` to the state of the graph after each node is called."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9d4d01da-7b64-4c92-96b7-29ec93332d0b",
      "metadata": {
        "id": "9d4d01da-7b64-4c92-96b7-29ec93332d0b",
        "outputId": "9f472f0c-2bf7-45c9-801a-8f2806be4b88"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--------------------Assistant Node--------------------\n",
            "{'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_SG7XYqDENGq7mwXrnioNLosS', 'function': {'arguments': '{\"a\":2,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-2c120fc3-3c82-4599-b8ec-24fbee207cad', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 2, 'b': 3}, 'id': 'call_SG7XYqDENGq7mwXrnioNLosS', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}\n",
            "--------------------Tools Node--------------------\n",
            "{'content': '6', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'tool', 'name': 'multiply', 'id': '3b40d091-58b2-4566-a84c-60af67206307', 'tool_call_id': 'call_SG7XYqDENGq7mwXrnioNLosS', 'artifact': None, 'status': 'success'}\n",
            "--------------------Assistant Node--------------------\n",
            "{'content': 'The result of multiplying 2 and 3 is 6.', 'additional_kwargs': {}, 'response_metadata': {'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_fde2829a40'}, 'type': 'ai', 'name': None, 'id': 'run-1272d9b0-a0aa-4ff7-8bad-fdffd27c5506', 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None}\n"
          ]
        }
      ],
      "source": [
        "initial_input = {\"messages\": HumanMessage(content=\"Multiply 2 and 3\")}\n",
        "thread = await client.threads.create()\n",
        "async for chunk in client.runs.stream(\n",
        "    thread[\"thread_id\"],\n",
        "    assistant_id = \"agent\",\n",
        "    input=initial_input,\n",
        "    stream_mode=\"updates\",\n",
        "):\n",
        "    if chunk.data:\n",
        "        assisant_node = chunk.data.get('assistant', {}).get('messages', [])\n",
        "        tool_node = chunk.data.get('tools', {}).get('messages', [])\n",
        "        if assisant_node:\n",
        "            print(\"-\" * 20+\"Assistant Node\"+\"-\" * 20)\n",
        "            print(assisant_node[-1])\n",
        "        elif tool_node:\n",
        "            print(\"-\" * 20+\"Tools Node\"+\"-\" * 20)\n",
        "            print(tool_node[-1])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8cc3bab2",
      "metadata": {
        "id": "8cc3bab2"
      },
      "source": [
        "Now, let's look at **replaying** from a specified checkpoint.\n",
        "\n",
        "We simply need to pass the `checkpoint_id`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d8ecc4fd",
      "metadata": {
        "id": "d8ecc4fd",
        "outputId": "d7b054c0-319b-4df2-e860-5b1009980f0d"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'values': {'messages': [{'content': 'Multiply 2 and 3',\n",
              "    'additional_kwargs': {'example': False,\n",
              "     'additional_kwargs': {},\n",
              "     'response_metadata': {}},\n",
              "    'response_metadata': {},\n",
              "    'type': 'human',\n",
              "    'name': None,\n",
              "    'id': 'df98147a-cb3d-4f1a-b7f7-1545c4b6f042',\n",
              "    'example': False}]},\n",
              " 'next': ['assistant'],\n",
              " 'tasks': [{'id': 'e497456f-827a-5027-87bd-b0ccd54aa89a',\n",
              "   'name': 'assistant',\n",
              "   'error': None,\n",
              "   'interrupts': [],\n",
              "   'state': None}],\n",
              " 'metadata': {'step': 0,\n",
              "  'run_id': '1ef6a449-7fbc-6c90-8754-4e6b1b582790',\n",
              "  'source': 'loop',\n",
              "  'writes': None,\n",
              "  'parents': {},\n",
              "  'user_id': '',\n",
              "  'graph_id': 'agent',\n",
              "  'thread_id': '708e1d8f-f7c8-4093-9bb4-999c4237cb4a',\n",
              "  'created_by': 'system',\n",
              "  'assistant_id': 'fe096781-5601-53d2-b2f6-0d3403f7e9ca'},\n",
              " 'created_at': '2024-09-03T22:33:51.380352+00:00',\n",
              " 'checkpoint_id': '1ef6a449-817f-6b55-8000-07c18fbdf7c8',\n",
              " 'parent_checkpoint_id': '1ef6a449-816c-6fd6-bfff-32a56dd2635f'}"
            ]
          },
          "execution_count": 24,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "states = await client.threads.get_history(thread['thread_id'])\n",
        "to_replay = states[-2]\n",
        "to_replay"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e33f865a",
      "metadata": {
        "id": "e33f865a"
      },
      "source": [
        "Let's stream with `stream_mode=\"values\"` to see the full state at every node as we replay."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "325e8272",
      "metadata": {
        "id": "325e8272",
        "outputId": "a6e6f7c6-9649-4e98-de98-711648641b93"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Receiving new event of type: metadata...\n",
            "{'run_id': '1ef6a44a-5806-6bb1-b2ee-92ecfda7f67d'}\n",
            "\n",
            "\n",
            "\n",
            "Receiving new event of type: values...\n",
            "{'messages': [{'content': 'Multiply 2 and 3', 'additional_kwargs': {'example': False, 'additional_kwargs': {}, 'response_metadata': {}}, 'response_metadata': {}, 'type': 'human', 'name': None, 'id': 'df98147a-cb3d-4f1a-b7f7-1545c4b6f042', 'example': False}]}\n",
            "\n",
            "\n",
            "\n",
            "Receiving new event of type: values...\n",
            "{'messages': [{'content': 'Multiply 2 and 3', 'additional_kwargs': {'example': False, 'additional_kwargs': {}, 'response_metadata': {}}, 'response_metadata': {}, 'type': 'human', 'name': None, 'id': 'df98147a-cb3d-4f1a-b7f7-1545c4b6f042', 'example': False}, {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'function': {'arguments': '{\"a\":2,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-e60d82d7-7743-4f13-bebd-3616a88720a9', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 2, 'b': 3}, 'id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}]}\n",
            "\n",
            "\n",
            "\n",
            "Receiving new event of type: values...\n",
            "{'messages': [{'content': 'Multiply 2 and 3', 'additional_kwargs': {'example': False, 'additional_kwargs': {}, 'response_metadata': {}}, 'response_metadata': {}, 'type': 'human', 'name': None, 'id': 'df98147a-cb3d-4f1a-b7f7-1545c4b6f042', 'example': False}, {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'function': {'arguments': '{\"a\":2,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-e60d82d7-7743-4f13-bebd-3616a88720a9', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 2, 'b': 3}, 'id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}, {'content': '6', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'tool', 'name': 'multiply', 'id': 'f1be0b83-4565-4aa2-9b9a-cd8874c6a2bc', 'tool_call_id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'artifact': None, 'status': 'success'}]}\n",
            "\n",
            "\n",
            "\n",
            "Receiving new event of type: values...\n",
            "{'messages': [{'content': 'Multiply 2 and 3', 'additional_kwargs': {'example': False, 'additional_kwargs': {}, 'response_metadata': {}}, 'response_metadata': {}, 'type': 'human', 'name': None, 'id': 'df98147a-cb3d-4f1a-b7f7-1545c4b6f042', 'example': False}, {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'function': {'arguments': '{\"a\":2,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-e60d82d7-7743-4f13-bebd-3616a88720a9', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 2, 'b': 3}, 'id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}, {'content': '6', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'tool', 'name': 'multiply', 'id': 'f1be0b83-4565-4aa2-9b9a-cd8874c6a2bc', 'tool_call_id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'artifact': None, 'status': 'success'}, {'content': 'The result of multiplying 2 and 3 is 6.', 'additional_kwargs': {}, 'response_metadata': {'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-55e5847a-d542-4977-84d7-24852e78b0a9', 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None}]}\n",
            "\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "async for chunk in client.runs.stream(\n",
        "    thread[\"thread_id\"],\n",
        "    assistant_id=\"agent\",\n",
        "    input=None,\n",
        "    stream_mode=\"values\",\n",
        "    checkpoint_id=to_replay['checkpoint_id']\n",
        "):\n",
        "    print(f\"Receiving new event of type: {chunk.event}...\")\n",
        "    print(chunk.data)\n",
        "    print(\"\\n\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "14c153b3",
      "metadata": {
        "id": "14c153b3"
      },
      "source": [
        "We can all view this as streaming only `updates` to state made by the nodes that we reply."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9e608e93",
      "metadata": {
        "id": "9e608e93",
        "outputId": "b02387af-9446-4417-f373-7fb77fcfaaf8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--------------------Assistant Node--------------------\n",
            "{'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_I2qudhMCwcw1GzcFN5q80rjj', 'function': {'arguments': '{\"a\":2,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-550e75ad-dbbc-4e55-9f00-aa896228914c', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 2, 'b': 3}, 'id': 'call_I2qudhMCwcw1GzcFN5q80rjj', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}\n",
            "--------------------Tools Node--------------------\n",
            "{'content': '6', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'tool', 'name': 'multiply', 'id': '731b7d4f-780d-4a8b-aec9-0d8b9c58c40a', 'tool_call_id': 'call_I2qudhMCwcw1GzcFN5q80rjj', 'artifact': None, 'status': 'success'}\n",
            "--------------------Assistant Node--------------------\n",
            "{'content': 'The result of multiplying 2 and 3 is 6.', 'additional_kwargs': {}, 'response_metadata': {'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-2326afa5-eb43-4568-b5ed-424c0a0fa076', 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None}\n"
          ]
        }
      ],
      "source": [
        "async for chunk in client.runs.stream(\n",
        "    thread[\"thread_id\"],\n",
        "    assistant_id=\"agent\",\n",
        "    input=None,\n",
        "    stream_mode=\"updates\",\n",
        "    checkpoint_id=to_replay['checkpoint_id']\n",
        "):\n",
        "    if chunk.data:\n",
        "        assisant_node = chunk.data.get('assistant', {}).get('messages', [])\n",
        "        tool_node = chunk.data.get('tools', {}).get('messages', [])\n",
        "        if assisant_node:\n",
        "            print(\"-\" * 20+\"Assistant Node\"+\"-\" * 20)\n",
        "            print(assisant_node[-1])\n",
        "        elif tool_node:\n",
        "            print(\"-\" * 20+\"Tools Node\"+\"-\" * 20)\n",
        "            print(tool_node[-1])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8e66e0e8",
      "metadata": {
        "id": "8e66e0e8"
      },
      "source": [
        "#### Forking\n",
        "\n",
        "Now, let's look at forking.\n",
        "\n",
        "Let's get the same step as we worked with above, the human input.\n",
        "\n",
        "Let's create a new thread with out agent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "01af5ed4",
      "metadata": {
        "id": "01af5ed4",
        "outputId": "55127dba-9b78-4edb-88e2-02dd6a2e64a1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--------------------Assistant Node--------------------\n",
            "{'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_HdWoyLELFZGEcqGxFt2fZzek', 'function': {'arguments': '{\"a\":2,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-cbd081b1-8cef-4ca8-9dd5-aceb134404dc', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 2, 'b': 3}, 'id': 'call_HdWoyLELFZGEcqGxFt2fZzek', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}\n",
            "--------------------Tools Node--------------------\n",
            "{'content': '6', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'tool', 'name': 'multiply', 'id': '11dd4a7f-0b6b-44da-b9a4-65f1677c8813', 'tool_call_id': 'call_HdWoyLELFZGEcqGxFt2fZzek', 'artifact': None, 'status': 'success'}\n",
            "--------------------Assistant Node--------------------\n",
            "{'content': 'The result of multiplying 2 and 3 is 6.', 'additional_kwargs': {}, 'response_metadata': {'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-936cf990-9302-45c7-9051-6ff1e2e9f316', 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None}\n"
          ]
        }
      ],
      "source": [
        "initial_input = {\"messages\": HumanMessage(content=\"Multiply 2 and 3\")}\n",
        "thread = await client.threads.create()\n",
        "async for chunk in client.runs.stream(\n",
        "    thread[\"thread_id\"],\n",
        "    assistant_id=\"agent\",\n",
        "    input=initial_input,\n",
        "    stream_mode=\"updates\",\n",
        "):\n",
        "    if chunk.data:\n",
        "        assisant_node = chunk.data.get('assistant', {}).get('messages', [])\n",
        "        tool_node = chunk.data.get('tools', {}).get('messages', [])\n",
        "        if assisant_node:\n",
        "            print(\"-\" * 20+\"Assistant Node\"+\"-\" * 20)\n",
        "            print(assisant_node[-1])\n",
        "        elif tool_node:\n",
        "            print(\"-\" * 20+\"Tools Node\"+\"-\" * 20)\n",
        "            print(tool_node[-1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3dbc8795-c3f5-4559-a00e-dc410c0a927f",
      "metadata": {
        "id": "3dbc8795-c3f5-4559-a00e-dc410c0a927f",
        "outputId": "3f255bb2-e05c-4c4c-ef49-ee9970736c0e"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'messages': [{'content': 'Multiply 2 and 3',\n",
              "   'additional_kwargs': {'example': False,\n",
              "    'additional_kwargs': {},\n",
              "    'response_metadata': {}},\n",
              "   'response_metadata': {},\n",
              "   'type': 'human',\n",
              "   'name': None,\n",
              "   'id': '93c18b95-9050-4a52-99b8-9374e98ee5db',\n",
              "   'example': False}]}"
            ]
          },
          "execution_count": 28,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "states = await client.threads.get_history(thread['thread_id'])\n",
        "to_fork = states[-2]\n",
        "to_fork['values']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "11e6cde1-0388-43ea-b994-1c4e9ca1199b",
      "metadata": {
        "id": "11e6cde1-0388-43ea-b994-1c4e9ca1199b",
        "outputId": "0f1eaca7-68b6-45ef-8632-2f6f649ddb42"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'93c18b95-9050-4a52-99b8-9374e98ee5db'"
            ]
          },
          "execution_count": 29,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "to_fork['values']['messages'][0]['id']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0c1e2300-c8b2-4994-a96d-1be19c04b6a8",
      "metadata": {
        "id": "0c1e2300-c8b2-4994-a96d-1be19c04b6a8",
        "outputId": "204d404a-0b0c-459e-8150-b8de24a652a2"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['assistant']"
            ]
          },
          "execution_count": 30,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "to_fork['next']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9d31d5aa-524f-42f4-ba7e-713a029610d6",
      "metadata": {
        "id": "9d31d5aa-524f-42f4-ba7e-713a029610d6",
        "outputId": "3b74df75-3f6a-4023-a892-2d10db5da2ad"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'1ef6a44b-27ec-681c-8000-ff7e345aee7e'"
            ]
          },
          "execution_count": 31,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "to_fork['checkpoint_id']"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8f11e1d9-9fe7-4243-a06f-9b07e38a12ad",
      "metadata": {
        "id": "8f11e1d9-9fe7-4243-a06f-9b07e38a12ad"
      },
      "source": [
        "Let's edit the state.\n",
        "\n",
        "Remember how our reducer on `messages` works:\n",
        "\n",
        "* It will append, unless we supply a message ID.\n",
        "* We supply the message ID to overwrite the message, rather than appending to state!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0198f1b8-2f57-4c6e-ac6a-c6fb80cce0bd",
      "metadata": {
        "id": "0198f1b8-2f57-4c6e-ac6a-c6fb80cce0bd"
      },
      "outputs": [],
      "source": [
        "forked_input = {\"messages\": HumanMessage(content=\"Multiply 3 and 3\",\n",
        "                                         id=to_fork['values']['messages'][0]['id'])}\n",
        "\n",
        "forked_config = await client.threads.update_state(\n",
        "    thread[\"thread_id\"],\n",
        "    forked_input,\n",
        "    checkpoint_id=to_fork['checkpoint_id']\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1dcd5b8e-6bb1-4967-84cf-4af710b8bf46",
      "metadata": {
        "id": "1dcd5b8e-6bb1-4967-84cf-4af710b8bf46",
        "outputId": "a0d74252-b8d6-4388-f56f-f45f88b7b050"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'configurable': {'thread_id': 'c99502e7-b0d7-473e-8295-1ad60e2b7ed2',\n",
              "  'checkpoint_ns': '',\n",
              "  'checkpoint_id': '1ef6a44b-90dc-68c8-8001-0c36898e0f34'},\n",
              " 'checkpoint_id': '1ef6a44b-90dc-68c8-8001-0c36898e0f34'}"
            ]
          },
          "execution_count": 33,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "forked_config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "015ac68a-5cc1-4c42-90a2-5b2b4865a153",
      "metadata": {
        "id": "015ac68a-5cc1-4c42-90a2-5b2b4865a153",
        "outputId": "a04edae2-9fd4-49ae-f375-68e78ae75c59"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'values': {'messages': [{'content': 'Multiply 3 and 3',\n",
              "    'additional_kwargs': {'additional_kwargs': {},\n",
              "     'response_metadata': {},\n",
              "     'example': False},\n",
              "    'response_metadata': {},\n",
              "    'type': 'human',\n",
              "    'name': None,\n",
              "    'id': '93c18b95-9050-4a52-99b8-9374e98ee5db',\n",
              "    'example': False}]},\n",
              " 'next': ['assistant'],\n",
              " 'tasks': [{'id': 'da5d6548-62ca-5e69-ba70-f6179b2743bd',\n",
              "   'name': 'assistant',\n",
              "   'error': None,\n",
              "   'interrupts': [],\n",
              "   'state': None}],\n",
              " 'metadata': {'step': 1,\n",
              "  'source': 'update',\n",
              "  'writes': {'__start__': {'messages': {'id': '93c18b95-9050-4a52-99b8-9374e98ee5db',\n",
              "     'name': None,\n",
              "     'type': 'human',\n",
              "     'content': 'Multiply 3 and 3',\n",
              "     'example': False,\n",
              "     'additional_kwargs': {},\n",
              "     'response_metadata': {}}}},\n",
              "  'parents': {},\n",
              "  'graph_id': 'agent'},\n",
              " 'created_at': '2024-09-03T22:34:46.678333+00:00',\n",
              " 'checkpoint_id': '1ef6a44b-90dc-68c8-8001-0c36898e0f34',\n",
              " 'parent_checkpoint_id': '1ef6a44b-27ec-681c-8000-ff7e345aee7e'}"
            ]
          },
          "execution_count": 34,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "states = await client.threads.get_history(thread['thread_id'])\n",
        "states[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3de80029-b987-49c5-890d-6cd70cbc8de7",
      "metadata": {
        "id": "3de80029-b987-49c5-890d-6cd70cbc8de7"
      },
      "source": [
        "To rerun, we pass in the `checkpoint_id`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "da005240-d3f0-4c89-9aca-dfcb5d410ceb",
      "metadata": {
        "id": "da005240-d3f0-4c89-9aca-dfcb5d410ceb",
        "outputId": "1856d759-bb5f-499d-958b-3f8d810c7589"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--------------------Assistant Node--------------------\n",
            "{'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_aodhCt5fWv33qVbO7Nsub9Q3', 'function': {'arguments': '{\"a\":3,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-e9759422-e537-4b9b-b583-36c688e13b4b', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 3, 'b': 3}, 'id': 'call_aodhCt5fWv33qVbO7Nsub9Q3', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}\n",
            "--------------------Tools Node--------------------\n",
            "{'content': '9', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'tool', 'name': 'multiply', 'id': '89787b0b-93de-4c0a-bea8-d2c3845534e1', 'tool_call_id': 'call_aodhCt5fWv33qVbO7Nsub9Q3', 'artifact': None, 'status': 'success'}\n",
            "--------------------Assistant Node--------------------\n",
            "{'content': 'The result of multiplying 3 by 3 is 9.', 'additional_kwargs': {}, 'response_metadata': {'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-0e16610f-4e8d-46f3-a5df-c2f187fae593', 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None}\n"
          ]
        }
      ],
      "source": [
        "async for chunk in client.runs.stream(\n",
        "    thread[\"thread_id\"],\n",
        "    assistant_id=\"agent\",\n",
        "    input=None,\n",
        "    stream_mode=\"updates\",\n",
        "    checkpoint_id=forked_config['checkpoint_id']\n",
        "):\n",
        "    if chunk.data:\n",
        "        assisant_node = chunk.data.get('assistant', {}).get('messages', [])\n",
        "        tool_node = chunk.data.get('tools', {}).get('messages', [])\n",
        "        if assisant_node:\n",
        "            print(\"-\" * 20+\"Assistant Node\"+\"-\" * 20)\n",
        "            print(assisant_node[-1])\n",
        "        elif tool_node:\n",
        "            print(\"-\" * 20+\"Tools Node\"+\"-\" * 20)\n",
        "            print(tool_node[-1])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "36956571-a2b8-4f1b-8e30-51f02f155a6f",
      "metadata": {
        "id": "36956571-a2b8-4f1b-8e30-51f02f155a6f"
      },
      "source": [
        "### LangGraph Studio\n",
        "\n",
        "Let's look at forking in the Studio UI with our `agent`, which uses `module-1/studio/agent.py` set in `module-1/studio/langgraph.json`."
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.1"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}